//
// Created by Administrator on 2021/6/11.
//


#ifndef DBNET_H
#define DBNET_H

#include "NvCaffeParser.h"
#include "NvInfer.h"
#include "NvOnnxParser.h"
#include "NvUffParser.h"
#include <opencv2/opencv.hpp>

#include "utils/cudaWrapper.h"
#include "utils/ioHelper.h"
#include "TextDetector/base_detect.h"

#define MAX_WORKSPACE_SIZE ((1ULL << 30) * 13) // 1 GB
#define MAX_BINDING_NUMS   (2)

class DBnet : public BaseDetect {
public:
    DBnet(int device_id, bool trt_serialize = false) : BaseDetect(device_id, trt_serialize) {};

    int init(const std::string &model_file);

    std::vector <std::vector<cv::Point2f>> extract(const cv::Mat &inputs);

    ~DBnet();

private:

    int serialization(const std::string &model_file, const std::string &engine_file);

    bool Input(const cv::Mat &src, float *data);

    std::vector <std::vector<cv::Point2f>> Output(std::vector<float> &output, int height, int width, int src_h, int src_w);

private:

    std::unique_ptr <nvinfer1::ICudaEngine, nvinfer1::Destroy<nvinfer1::ICudaEngine>> mEngine{nullptr};
    std::unique_ptr <nvinfer1::IExecutionContext, nvinfer1::Destroy<nvinfer1::IExecutionContext>> context{nullptr};

    void *bindings[MAX_BINDING_NUMS]{0};
    cudaStream_t stream_t;

    cv::Mat input;

    //int device_id = 0;
    int channels = 3;
    int inputH, inputW;
    int limit_side_len = 1920;
    //bool trt_serialize;

};

#endif //DBNET_H
